import socket
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
from Crypto.Hash import SHA256
import rsa


# 解密文件
def decrypt_file(encrypted_file_path, key, iv):
    cipher = AES.new(key, AES.MODE_CBC, iv)

    with open(encrypted_file_path, 'rb') as f:
        ciphertext = f.read()

    decrypted_data = unpad(cipher.decrypt(ciphertext), AES.block_size)

    decrypted_file_path = encrypted_file_path.replace('.enc', '.txt')
    # decrypted_file_path = encrypted_file_path.replace('.enc', '.zip')
    with open(decrypted_file_path, 'wb') as f:
        f.write(decrypted_data)

    print(f"文件已解密并保存为: {decrypted_file_path}")


# 验证签名
def verify_signature(file_path, signature_path, public_key_path):
    # file_path = 'file_to_send.txt'
    # signature_path = 'signature.sig'
    with open(file_path, 'rb') as f:
        file_data = f.read()

    h = SHA256.new(file_data)

    with open(public_key_path, 'rb') as pub_file:
        pub_key = rsa.PublicKey.load_pkcs1(pub_file.read())

    with open(signature_path, 'rb') as sig_file:
        signature = sig_file.read()
    try:
        rsa.verify(h.digest(), signature, pub_key)
        print("签名验证成功！")
        return True
    except rsa.VerificationError:
        print("签名验证失败！")
        return False


# 服务器设置
def start_server(host, port, key, iv, public_key_path):
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind((host, port))
    server_socket.listen(1)
    print(f"服务器启动，监听端口 {port}...")

    conn, addr = server_socket.accept()
    print(f"连接来自: {addr}")

    # 接收文件的大小
    file_size = int(conn.recv(16).decode())
    print(f"接收到的文件大小: {file_size} 字节")

    # 接收加密的文件数据
    with open("received.enc", 'wb') as f:
        remaining_size = file_size
        while remaining_size > 0:
            data = conn.recv(1024)
            f.write(data)
            remaining_size -= len(data)

    # 接收签名
    signature_size = int(conn.recv(256).decode())
    # signature = conn.recv(signature_size)
    with open("received_signature.sig", 'wb') as f:
        remaining_size = signature_size
        while remaining_size > 0:
            data = conn.recv(256)
            f.write(data)
            remaining_size -= len(data)

    print("文件接收完成，正在验证签名...")
    decrypt_file('received.enc', key, iv)
    if verify_signature("received.txt", "received_signature.sig", public_key_path):
        print("文件验证通过")
    else:
        print("文件验证失败！")

    conn.close()
    server_socket.close()


if __name__ == "__main__":
    # 设置 AES 加密的密钥和 IV
    # key = b'1234567890123456'  # 16字节的密钥
    with open('shared_secret_A_path.txt', 'rb') as key_file:
        key = key_file.read()
    key = key + b'\x00' * (16 - len(key))
    iv = b'1234567890123456'  # 16字节的初始化向量
    public_key_path = "public.pem"

    start_server('0.0.0.0', 12345, key, iv, public_key_path)
    # verify_signature(0, 0, public_key_path)
